from pymongo import MongoClient, ASCENDING, TEXT
from datetime import datetime

class DatabaseInitializer:
    def __init__(self):
        self.client = MongoClient('mongodb://zsl:zsl201115@localhost:27017/medical_regulatory_kb')
        self.db = self.client['medical_regulatory_kb']
    
    def create_collections(self):
        """Create all necessary collections with indexes"""
        
        # 1. Documents Collection
        if 'documents' not in self.db.list_collection_names():
            self.db.create_collection('documents')
        
        documents = self.db.documents
        documents.create_index([('filename', ASCENDING)], unique=True)
        documents.create_index([('file_name', ASCENDING)])  # Index for new file_name field
        documents.create_index([('country', ASCENDING)])
        documents.create_index([('document_type', ASCENDING)])
        documents.create_index([('language', ASCENDING)])  # Keep for potential future use
        documents.create_index([('regulation_id', ASCENDING)])
        documents.create_index([('issuing_authority', ASCENDING)])  # Index for issuing authority
        documents.create_index([('collection', ASCENDING)])
        documents.create_index([('tags', ASCENDING)])  # Index for tags array
        documents.create_index([('status', ASCENDING)])  # Index for status field
        documents.create_index([('update_time', ASCENDING)])  # Index for update_time field
        documents.create_index([('content_text', TEXT)])
        
        # 2. Document Chunks Collection (for vector embeddings)
        if 'document_chunks' not in self.db.list_collection_names():
            self.db.create_collection('document_chunks')
        
        chunks = self.db.document_chunks
        chunks.create_index([('document_id', ASCENDING)])
        chunks.create_index([('language', ASCENDING)])
        chunks.create_index([('collection', ASCENDING)]) 
        
        # 3. QA Sessions Collection
        if 'qa_sessions' not in self.db.list_collection_names():
            self.db.create_collection('qa_sessions')
        
        qa_sessions = self.db.qa_sessions
        qa_sessions.create_index([('session_id', ASCENDING)])
        qa_sessions.create_index([('user_id', ASCENDING)])
        qa_sessions.create_index([('timestamp', ASCENDING)])
        
        # 4. Collections Metadata Collection 
        if 'collections' not in self.db.list_collection_names():
            self.db.create_collection('collections')
        
        collections = self.db.collections
        collections.create_index([('name', ASCENDING)], unique=True)
        
        # 5. Tags Collection (NEW)
        if 'tags' not in self.db.list_collection_names():
            self.db.create_collection('tags')
        
        tags = self.db.tags
        tags.create_index([('name', ASCENDING)], unique=True)
        tags.create_index([('usage_count', ASCENDING)])

        # Insert default collection data
        self._create_default_collections()

        print("✓ Database collections created successfully!")

        # 6. Collection Info Collection
        if 'collection_info' not in self.db.list_collection_names():
            self.db.create_collection('collection_info')
        
        collection_info = self.db.collection_info
        collection_info.create_index([('name', ASCENDING)], unique=True)
        collection_info.create_index([('display_name', ASCENDING)])
        collection_info.create_index([('country', ASCENDING)])
        collection_info.create_index([('authority', ASCENDING)])
        
        print("✓ Collection info collection created successfully!")
    

    def _create_default_collections(self):
        """Create default regulatory authority collections"""
        default_collections = [
            {
                'name': 'FDA',
                'display_name': 'US FDA',
                'description': 'US FDA regulations, guidances, and submission requirements for medical device compliance.',
                'country': 'USA',
                'authority': 'FDA',
                'document_count': 0,
                'color': 'blue',
                'created_at': datetime.now()
            },
            {
                'name': 'NMPA',
                'display_name': 'China NMPA', 
                'description': 'NMPA requirements and registration documentation for medical devices in China.',
                'country': 'China',
                'authority': 'NMPA',
                'document_count': 0,
                'color': 'red',
                'created_at': datetime.now()
            },
            {
                'name': 'PMDA',
                'display_name': 'Japan PMDA',
                'description': 'PMDA approval and compliance documentation for medical devices in Japan.',
                'country': 'Japan', 
                'authority': 'PMDA',
                'document_count': 0,
                'color': 'green',
                'created_at': datetime.now()
            },
            {
                'name': 'ANVISA',
                'display_name': 'Brazil ANVISA',
                'description': 'ANVISA registration materials and regulations for medical devices in Brazil.',
                'country': 'Brazil',
                'authority': 'ANVISA', 
                'document_count': 0,
                'color': 'purple',
                'created_at': datetime.now()
            }
        ]
        
        collections = self.db.collections
        collection_info = self.db.collection_info

        for collection in default_collections:
            # 插入到原有的collections集合
            if not collections.find_one({'name': collection['name']}):
                collections.insert_one(collection)
            
            # 同时插入到新的collection_info集合
            if not collection_info.find_one({'name': collection['name']}):
                collection_info.insert_one(collection)
    
    def verify_setup(self):
        """Verify database setup"""
        collections = self.db.list_collection_names()
        print(f"\nAvailable collections: {collections}")
        
        for collection_name in collections:
            indexes = self.db[collection_name].index_information()
            print(f"\n{collection_name} indexes: {list(indexes.keys())}")

if __name__ == "__main__":
    initializer = DatabaseInitializer()
    initializer.create_collections()
    initializer.verify_setup()